#include "kalloc.h"
#include "printf.h" 
#include "console.h"
#include "types.h"
#include "string.h"

typedef struct { 
    volatile int locked; 
} spinlock_t;

static inline void spinlock_init(spinlock_t *lk) { lk->locked = 0; }
static inline void spinlock_acquire(spinlock_t *lk) {
    while (__sync_lock_test_and_set(&lk->locked, 1)) {
        while (lk->locked) { /* busy wait */ }
    }
}
static inline void spinlock_release(spinlock_t *lk) {
    __sync_lock_release(&lk->locked);
}

// 内存管理常量
#define PGSIZE 4096
#define PGROUNDUP(sz) (((sz) + PGSIZE - 1) & ~(PGSIZE - 1))
#define PGROUNDDOWN(a) ((a) & ~(PGSIZE - 1))

// 内存范围
#define DRAM_BASE 0x80000000UL
#define PHYSTOP   0x88000000UL  // 指定128MB，可修改

// 空闲链表节点
struct run {
    struct run *next;
};

// 内存管理器全局状态
static struct {
    spinlock_t lock;
    struct run *freelist;
    uint64_t total_pages;
    uint64_t free_pages;
} kmem;

// 释放一段内存范围
static void freerange(void* pa_start, void* pa_end) {
    char* p = (char*)PGROUNDUP((uint64_t)pa_start);
    for (; p + PGSIZE <= (char*)pa_end; p += PGSIZE) {
        kfree(p);
    }
}

// 初始化物理内存管理器
void kinit(void) {
    spinlock_init(&kmem.lock);
    
    // 获取内核结束地址
    extern char _end[];
    uint64_t start = PGROUNDUP((uint64_t)_end);
    uint64_t end = PHYSTOP;
    
    kmem.total_pages = (end - start) / PGSIZE;
    kmem.free_pages = 0;
    
    cprintf("KALLOC: initializing physical memory manager\n");
    cprintf("KALLOC: memory range [0x%lx, 0x%lx), total pages: %lu\n", 
           start, end, kmem.total_pages);
    
    freerange((void*)start, (void*)end);
    cprintf("KALLOC: initialization completed, free pages: %lu\n", kmem.free_pages);
}

// 分配单个物理页
void* kalloc(void) {
    spinlock_acquire(&kmem.lock);
    
    if (kmem.freelist == NULL) {
        spinlock_release(&kmem.lock);
        cprintf_color(31, "KALLOC: out of physical memory!\n");
        return NULL;
    }
    
    struct run* r = kmem.freelist;
    kmem.freelist = r->next;
    kmem.free_pages--;
    
    spinlock_release(&kmem.lock);
    
    // 清空页面内容（安全）
    memset(r, 0, PGSIZE);
    
    return (void*)r;
}

// 释放单个物理页
void kfree(void* pa) {
    // 基本安全检查
    if ((uint64_t)pa % PGSIZE != 0 || (uint64_t)pa < DRAM_BASE) {
        cprintf_color(31, "KALLOC: invalid page free attempt: 0x%lx\n", (uint64_t)pa);
        return;
    }

    spinlock_acquire(&kmem.lock);

    // 简单的double-free检测
    struct run* r = (struct run*)pa;
    for (struct run* p = kmem.freelist; p; p = p->next) {
        if (p == r) {
            spinlock_release(&kmem.lock);
            cprintf_color(31, "KALLOC: double-free detected at 0x%lx\n", (uint64_t)pa);
            return;
        }
    }

    // 用特殊模式填充页面（帮助检测use-after-free）
    //memset(pa, 0xAA, PGSIZE);

    // 插入空闲链表头部
    r->next = kmem.freelist;
    kmem.freelist = r;
    kmem.free_pages++;
    
    spinlock_release(&kmem.lock);
}

// 分配连续的n个物理页（简化版）
void* kalloc_pages(int n) {
    if (n <= 0) return NULL;
    
    // 简单实现：多次分配单页（不保证连续性）
    void* result = kalloc();
    if (result == NULL) return NULL;
    
    return result;
}

// 获取空闲页数
uint64_t get_free_page_count(void) {
    return kmem.free_pages;
}

// 获取总页数
uint64_t get_total_page_count(void) {
    return kmem.total_pages;
}
